import math
from scipy import constants as const
import numpy as np

Te = 20.0
Ti = 20.0

ne = 1.0e19
ni = 1.0e19

B_mag = 2.0
B_angle_alpha = 10.0 * math.pi / 180.0
B_x = B_mag * math.sin(B_angle_alpha)
B_y = B_mag * math.cos(B_angle_alpha)
B_z = 0.0

dx = 5.0e-5


# mass of H: 1.67262158e-27, D: 2.0 * 1.67262158e-27, T: 3.0 * 1.67262158e-27
me = 9.109382616e-31
mi = 2.0 * 1.67262158e-27
mr = me * mi / ( me + mi )
qi = 1

debye_length_e      = math.sqrt(const.epsilon_0 * Te / (ne * const.e))
debye_length_ion    = math.sqrt(const.epsilon_0 * Ti / (ni * qi * const.e))
debye_length        = math.sqrt(const.epsilon_0 * Te *  Ti / ((Te + Ti) * ne * qi * const.e))
period_plasma       = math.sqrt(const.epsilon_0 * me / (ne * const.e * const.e)) * 2.0 * const.pi

ion_sound_speed     = math.sqrt(Te * const.e / mi)
thermal_speed_e     = math.sqrt(Te * const.e / me)
thermal_speed_ion   = math.sqrt(Ti * const.e / mi)
ve                  = thermal_speed_e
vi                  = thermal_speed_ion

ion_sound_speed_x = ion_sound_speed * math.sin(B_angle_alpha)
ion_sound_speed_y = ion_sound_speed * math.cos(B_angle_alpha)
ion_sound_speed_z = 0.0



Omega_e             = const.e * B_mag / me
Omega_i             = const.e * B_mag / mi
rotation_period_e   = 2.0 * const.pi * me / (const.e * B_mag)
rotation_period_ion = 2.0 * const.pi * mi / (qi * const.e * B_mag)
rotation_radius_e   = me * thermal_speed_e / (const.e * B_mag)
rotation_radius_ion = mi * thermal_speed_ion / (qi * const.e * B_mag)

particle_flux          = ni * ion_sound_speed * math.sin(B_angle_alpha)
ion_saturation_current = ni * ion_sound_speed * const.e
macro_particle_weight  = ne * dx * dx * dx / 100.0

#================== diagnostic ne Te Is ========================
#saturated ion current density A/m^2

Js = np.array([6.364, 5.271, 20.81, 15.47])
Js = Js  * 1.0e4
#divertor temperature and density: tet, net
tet = np.array([1.746, 3.238, 10.93, 5.117]) * 5.0
miH = mi  / 2.0
net = Js * 1.0 / (const.e * np.sqrt(const.e * tet / miH))

print("==========================================")
print("debye_length:        ", debye_length)
print("debye_length_e:      ", debye_length_e)
print("debye_length_ion:    ", debye_length_ion)
print("period_plasma:       ", period_plasma)
print("==========================================")
print(" ")

print("==========================================")
print("rotation_radius_e:    ", rotation_radius_e)
print("rotation_radius_ion:  ", rotation_radius_ion)
print("rotation_period_e:    ", rotation_period_e)
print("rotation_period_ion:  ", rotation_period_ion)
print("==========================================")
print(" ")

print("==========================================")
print("magnetic field:         ", B_x, ",", B_y, ",", B_z)
print("ion_sound_speed:        ", ion_sound_speed, ion_sound_speed_x, ",", ion_sound_speed_y, ",", ion_sound_speed_z)
print("thermal_speed_ion:      ", thermal_speed_ion)
print("particle_flux:          ", particle_flux)
print("ion_saturation_current: ", ion_saturation_current)
print("macro_particle_weight:  ", macro_particle_weight)
print("==========================================")
print(" ")



print("================== diagnostic ne Te Is ========================")
print("divertor density:  ", net)
print("========================================================================")
print(" ")

